import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.utils import initialize_weights
import numpy as np
import math
from models.model_clam import CLAM_SB
from models.Resnet1d import ResNet, BasicBlock, conv3x1
import torch.utils.data as Data
from torch.nn import init
from models.model_attn import Attn_Net, Attn_Net_Gated, Attn_External, \
    Attn_ConvNet, Attn_CoordNet
from models.blocks import BatchMinMaxNorm, BatchMeanNorm, ChannelBlock
from models.loss_function import LGMLoss, LGMLoss_v0
from pytorch_metric_learning import losses


class CLAM_TB(CLAM_SB):
    def __init__(self, size_arg="small", dropout=False, k_sample=8, n_classes=2,
                 instance_loss_fn=nn.CrossEntropyLoss(), subtyping=False):
        super(CLAM_TB, self).__init__()
        self.size_dict = {"small": [1024, 512, 256], "big": [1024, 512, 384]}
        size = self.size_dict[size_arg]

        fc1 = [nn.Linear(size[0], size[1]), nn.ReLU()]
        if dropout:
            fc1.append(nn.Dropout(0.25))
        self.fc1 = nn.Sequential(*fc1)

        attn_type = 5
        if attn_type == 1:
            self.attention_net = Attn_Net_Gated(L=size[1], D=size[2], dropout=dropout, n_classes=1)
        elif attn_type == 2:
            self.attention_net = Attn_Net(L=size[1], D=size[2], dropout=dropout, n_classes=1)
        elif attn_type == 3:
            self.attention_net = Attn_ConvNet(L=512, D=256, n_classes=1)
        elif attn_type == 4:
            self.attention_net = Attn_External(d_model=512, S=64, n_classes=1)
        elif attn_type == 5:
            self.attention_net = Attn_CoordNet(L=512, D=256, n_classes=1)
        else:
            raise NotImplementedError

        self.classifiers = nn.Sequential(nn.Linear(size[1], n_classes))
        instance_classifiers = [nn.Linear(size[1], 2) for i in range(n_classes)]
        self.instance_classifiers = nn.ModuleList(instance_classifiers)
        self.k_sample = k_sample
        self.instance_loss_fn = instance_loss_fn
        self.n_classes = n_classes
        self.subtyping = subtyping

        initialize_weights(self)

    def relocate(self):
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.fc1 = self.fc1.to(device)
        self.attention_net = self.attention_net.to(device)
        self.classifiers = self.classifiers.to(device)
        self.instance_classifiers = self.instance_classifiers.to(device)

    def forward(self, h, label=None, coords=None, instance_eval=False, return_features=False, attention_only=False):
        device = h.device
        h = self.fc1(h)
        A, h = self.attention_net(h, coords)  # NxK

        A = torch.transpose(A, 1, 0)  # KxN
        if attention_only:
            return A
        A_raw = A
        A = F.softmax(A, dim=1)  # softmax over N，找出显著的样本

        if instance_eval:
            total_inst_loss = 0.0
            all_preds = []
            all_targets = []
            inst_labels = F.one_hot(label, num_classes=self.n_classes).squeeze()  # binarize label
            for i in range(len(self.instance_classifiers)):
                inst_label = inst_labels[i].item()
                classifier = self.instance_classifiers[i]
                if inst_label == 1:  # in-the-class:
                    instance_loss, preds, targets = self.inst_eval(A, h, classifier)
                    all_preds.extend(preds.cpu().numpy())
                    all_targets.extend(targets.cpu().numpy())
                else:  # out-of-the-class
                    if self.subtyping:
                        instance_loss, preds, targets = self.inst_eval_out(A, h, classifier)
                        all_preds.extend(preds.cpu().numpy())
                        all_targets.extend(targets.cpu().numpy())
                    else:
                        continue
                total_inst_loss += instance_loss

            if self.subtyping:
                total_inst_loss /= len(self.instance_classifiers)

        M = torch.mm(A, h)

        logits = self.classifiers(M)
        Y_hat = torch.topk(logits, k=1, dim=1)[1]  # 沿给定dim维度返回输入张量logits中 k 个最大值。
        Y_prob = F.softmax(logits, dim=1)
        if instance_eval:
            results_dict = {'instance_loss': total_inst_loss, 'inst_labels': np.array(all_targets),
                            'inst_preds': np.array(all_preds)}
        else:
            results_dict = {}
        if return_features:
            results_dict.update({'features': M})
        return logits, Y_prob, Y_hat, A_raw, results_dict



       